import numpy as np
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
from matplotlib.patches import Rectangle
from PyQt5.QtWidgets import QApplication, QMainWindow
import sys


class SelectBox:
    def __init__(self, ax):
        self.ax = ax
        self.rect = None
        self.press = None

    def on_press(self, event):
        if event.inaxes != self.ax:
            return
        self.press = (event.xdata, event.ydata)

    def on_release(self, event):
        if event.inaxes != self.ax:
            return
        if self.press is None:
            return

        x0, y0 = self.press
        x1, y1 = event.xdata, event.ydata

        # 框选完成，可以执行后续操作
        print(f"Selected area: ({x0}, {y0}) to ({x1}, {y1})")

        # 如果需要，可以移除矩形以准备下一次框选
        if self.rect is not None:
            self.rect.remove()
            self.rect = None

        self.ax.figure.canvas.draw_idle()

    def on_motion(self, event):
        if self.press is None:
            return
        if event.inaxes != self.ax:
            return

        x0, y0 = self.press
        x1, y1 = event.xdata, event.ydata

        # 确保框选是矩形
        x0, x1 = min(x0, x1), max(x0, x1)
        y0, y1 = min(y0, y1), max(y0, y1)

        # 如果之前已经有一个矩形，先移除它
        if self.rect is not None:
            self.rect.remove()

            # 创建一个新的矩形并添加到axes中
        self.rect = Rectangle((x0, y0), x1 - x0, y1 - y0, edgecolor='r', facecolor='none', linewidth=2)
        self.ax.add_patch(self.rect)

        self.ax.figure.canvas.draw_idle()


class DualAxisWindow(QMainWindow):
    def __init__(self):
        super().__init__()

        self.fig = Figure()
        self.ax = self.fig.add_subplot(111)

        # 假设有一些数据
        self.ax.plot(np.random.rand(10), np.random.rand(10), 'o')

        self.select_box = SelectBox(self.ax)

        self.canvas = FigureCanvas(self.fig)
        self.setCentralWidget(self.canvas)

        # 绑定事件
        self.canvas.mpl_connect('button_press_event', self.select_box.on_press)
        self.canvas.mpl_connect('button_release_event', self.select_box.on_release)
        self.canvas.mpl_connect('motion_notify_event', self.select_box.on_motion)


if __name__ == '__main__':
    app = QApplication([])
    mainWin = DualAxisWindow()
    mainWin.show()
    sys.exit(app.exec_())